package jwt

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"strings"
)

/**
 * JWT (JSON Web Tokens) is a compact URL-safe means of representing claims to be transferred between two parties.
 * It is a standard for securely transmitting information between a server and a client.
 *
 * JWT consists of three parts:
 * 1. Header: Contains the type and the hashing algorithm used.
 * 2. Payload: Contains the claims (data) to be transmitted.
 * 3. Signature: A cryptographic hash of the header and payload using the secret key.
 *
 * JWTs are typically used for authentication and authorization purposes.
 *
 * Example usage:
 *
    // Step 1: Generate a JWT
    claims := &Claims{
        UserID: "user123",
        Name:   "John Doe",
        Email:  "john.doe@example.com",
    }
    token, err := GenerateJWT(claims, "secret_key")

	jwt : Json Web Token
	由三部分构成：
	header：alg：表示 hash 算法，typ：JWT
	payload：sub：表示 subject，exp：表示过期时间，iat：表示创建时间
	signature：header + payload + secret_key 经过 Base64 编码后得到
*/

var (
	DefaultHeader = JwtHeader{
		Alg: "HS256",
		Typ: "JWT",
	}

	UID_IN_TOKEN = "uuid"
)

type JwtHeader struct {
	Alg string `json:"alg"`
	Typ string `json:"typ"`
}

type JwtPayload struct {
	ID          string         `json:"jti"`
	Issue       string         `json:"iss"`
	Audience    string         `json: "aud"`
	Subject     string         `json:"sub"`
	IssueAt     int64          `json:"iat"`
	NotBefore   int64          `json:"nbf"`
	Expiration  int64          `json:"exp"`
	UserDefined map[string]any `json:"ud"`
}

func GenJwt(header JwtHeader, payload JwtPayload, secret string) (string, error) {
	var part1, part2, signature string
	if bs1, err := json.Marshal(header); err != nil {
		return "", err
	} else {
		part1 = base64.RawURLEncoding.EncodeToString(bs1)
	}
	if bs2, err := json.Marshal(payload); err != nil {
		return "", err
	} else {
		part2 = base64.RawURLEncoding.EncodeToString(bs2)
	}
	h := hmac.New(sha256.New, []byte(secret))
	h.Write([]byte(part1 + "." + part2))
	signature = base64.RawURLEncoding.EncodeToString(h.Sum(nil))
	jwt := part1 + "." + part2 + "." + signature
	return jwt, nil
}

func VerifyJwt(token string, secret string) (*JwtHeader, *JwtPayload, error) {
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return nil, nil, fmt.Errorf("invalid token")
	}

	h := hmac.New(sha256.New, []byte(secret))
	h.Write([]byte(parts[0] + "." + parts[1]))
	signature := base64.RawURLEncoding.EncodeToString(h.Sum(nil))
	if signature != parts[2] {
		return nil, nil, fmt.Errorf("invalid signature")
	}
	var part1, part2 []byte
	var err error
	if part1, err = base64.RawURLEncoding.DecodeString(parts[0]); err != nil {
		return nil, nil, err
	}
	if part2, err = base64.RawURLEncoding.DecodeString(parts[1]); err != nil {
		return nil, nil, err
	}
	var header JwtHeader
	if err := json.Unmarshal(part1, &header); err != nil {
		return nil, nil, err
	}
	var payload JwtPayload
	if err := json.Unmarshal(part2, &payload); err != nil {
		return nil, nil, err
	}
	return &header, &payload, nil
}
